"""
朴素贝叶斯
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import sklearn.naive_bayes as nb
import sklearn.model_selection as ms
import sklearn.metrics as sm


# 加载数据
data = pd.read_csv('multiple1.txt',header=None,names=['x1','x2','y'])
plt.scatter(data['x1'],data['x2'],c=data['y'],cmap='brg')
# plt.show()

# 整理数据
x = data.iloc[:,:-1]
y = data.iloc[:,-1]
# 划分训练集和测试集
train_x,test_x,train_y,test_y = ms.train_test_split(x,y,test_size=0.1,random_state=7)
# 建立模型
# 高斯朴素贝叶斯分类器
model = nb.GaussianNB()
model.fit(train_x,train_y)
pred_test_y = model.predict(test_x)
print(sm.classification_report(test_y,pred_test_y))

# 画出分类边界线
x1s = np.linspace(data['x1'].min(),data['x1'].max(),200)
x2s = np.linspace(data['x2'].min(),data['x2'].max(),200)

points = []
for x1 in x1s:
    for x2 in x2s:
        points.append([x1,x2])
points = pd.DataFrame(points,columns=['x1','x2'])

points_label = model.predict(points)
plt.scatter(points['x1'],points['x2'],c=points_label,cmap='gray')
plt.scatter(data['x1'],data['x2'],c=data['y'],cmap='brg')
plt.show()